# -*- coding: utf-8 -*-
"""
@author: Trae
@contact: traetai@gmail.com
@software: TraeAI
@file: visualization.py
@time: 2024/7/29 16:45
@desc: 结果可视化工具
"""
import cv2
import numpy as np

# 为了美观，定义一些颜色
COLORS = np.random.uniform(0, 255, size=(100, 3))

def draw_detections(frame, detections):
    """
    在图像帧上绘制检测结果。

    Args:
        frame: 待绘制的图像帧。
        detections: YOLOv8的检测结果对象。

    Returns:
        绘制了检测框的图像帧。
    """
    if detections is None or len(detections.boxes) == 0:
        return frame

    # 获取边界框、置信度和类别ID
    boxes = detections.boxes.xyxy.cpu().numpy()  # .xyxy格式
    confidences = detections.boxes.conf.cpu().numpy()
    class_ids = detections.boxes.cls.cpu().numpy().astype(int)
    class_names = detections.names  # 获取类别名称字典

    for i in range(len(boxes)):
        x1, y1, x2, y2 = map(int, boxes[i])
        confidence = confidences[i]
        class_id = class_ids[i]
        class_name = class_names.get(class_id, 'Unknown')
        color = COLORS[class_id % len(COLORS)]

        # 绘制边界框
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

        # 准备标签文本
        label = f'{class_name}: {confidence:.2f}'

        # 计算文本尺寸以绘制背景
        (label_width, label_height), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
        cv2.rectangle(frame, (x1, y1 - label_height - baseline), (x1 + label_width, y1), color, -1)

        # 绘制文本
        cv2.putText(frame, label, (x1, y1 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    return frame

if __name__ == '__main__':
    # 这是一个测试该模块功能的示例
    # 依赖于detector.py，需要先运行它以理解其输出
    import sys
    from pathlib import Path
    # 将src目录添加到sys.path
    sys.path.append(str(Path(__file__).resolve().parents[1]))

    from core.detector import Detector

    # 假设项目根目录下有yolov8s/yolov8s.pt和video/test.jpg
    root_dir = Path(__file__).resolve().parents[2]
    model_path = root_dir / 'yolov8s' / 'yolov8s.pt'
    image_path = root_dir / 'video' / 'test.jpg' # 你需要一张名为test.jpg的测试图片

    if not model_path.exists() or not image_path.exists():
        print(f"Error: Ensure model '{model_path}' and image '{image_path}' exist.")
    else:
        detector = Detector(str(model_path))
        img = cv2.imread(str(image_path))
        detections = detector.detect(img, target_classes=[0, 2]) # 检测人和车

        # 使用我们自己的绘制函数
        annotated_frame = draw_detections(img.copy(), detections)

        cv2.imshow("Custom Detections Visualization", annotated_frame)
        cv2.waitKey(0)
        cv2.destroyAllWindows()